import os
import threading
from PIL import Image, ImageOps
import numpy as np
from PySide6.QtCore import Signal, QObject

class ProcessorSignals(QObject):
    """图像处理器信号类，用于在处理过程中发射进度信号"""
    
    progress_changed = Signal(int)  # 进度信号(0-100)
    processing_finished = Signal(object)  # 处理完成信号，传递处理后的图像
    processing_error = Signal(str)  # 处理错误信号，传递错误信息
    batch_progress_changed = Signal(int, int)  # 批处理进度信号(当前索引, 总数)
    batch_processing_finished = Signal(list)  # 批处理完成信号，传递处理后的图像列表
    batch_processing_error = Signal(str, int)  # 批处理错误信号，传递错误信息和索引

class ImageProcessor:
    """图像处理类，用于将彩色BMP图片转换为黑白像素图像，并支持导出为BMP或C文件格式"""
    
    def __init__(self):
        self.original_image = None
        self.processed_image = None
        self.batch_images = []  # 批处理图像列表
        self.batch_processed_images = []  # 批处理后的图像列表
        self.r_threshold = 128  # 默认红色阈值128
        self.g_threshold = 128  # 默认绿色阈值128
        self.b_threshold = 128  # 默认蓝色阈值128
        self.invert_colors = False  # 默认不反转颜色
        
        # 预处理选项
        self.resize_enabled = False
        self.resize_width = 128
        self.resize_height = 64
        self.rotate_enabled = False
        self.rotate_angle = 0
        self.flip_horizontal = False
        self.flip_vertical = False
        
        self.signals = ProcessorSignals()  # 信号对象
        self.processing_thread = None  # 处理线程
        self.is_processing = False  # 处理状态标志
        
    def load_image(self, image_path):
        """加载图像文件

        Args:
            image_path (str): 图像文件路径
            
        Returns:
            bool: 加载成功返回True，否则返回False
        """
        try:
            if not os.path.exists(image_path):
                return False
                
            # 使用PIL打开图像
            self.original_image = Image.open(image_path)
            
            # 确保图像是BMP格式
            if self.original_image.format != 'BMP':
                self.original_image = self.original_image.convert('RGB')
            
            # 转换完成后设置processed_image为None，表示需要重新处理
            self.processed_image = None
            return True
        except Exception as e:
            print(f"加载图像失败: {str(e)}")
            return False
            
    def load_images(self, image_paths):
        """加载多个图像文件用于批处理
        
        Args:
            image_paths (list): 图像文件路径列表
            
        Returns:
            bool: 如果至少一个图像加载成功返回True，否则返回False
        """
        self.batch_images = []
        success = False
        
        for path in image_paths:
            try:
                if not os.path.exists(path):
                    continue
                    
                # 使用PIL打开图像
                img = Image.open(path)
                
                # 确保图像是BMP格式
                if img.format != 'BMP':
                    img = img.convert('RGB')
                
                # 添加到批处理列表
                self.batch_images.append({"path": path, "image": img})
                success = True
            except Exception as e:
                print(f"加载图像失败 {path}: {str(e)}")
        
        # 清空批处理结果列表
        self.batch_processed_images = []
        
        return success
    
    def set_thresholds(self, r_threshold, g_threshold, b_threshold):
        """设置RGB颜色阈值

        Args:
            r_threshold (int): 红色阈值(0-255)
            g_threshold (int): 绿色阈值(0-255)
            b_threshold (int): 蓝色阈值(0-255)
        """
        self.r_threshold = r_threshold
        self.g_threshold = g_threshold
        self.b_threshold = b_threshold
        # 重置processed_image，需要重新处理
        self.processed_image = None
        
    def set_invert_colors(self, invert):
        """设置是否反转颜色
        
        Args:
            invert (bool): True表示反转颜色，False表示不反转
        """
        self.invert_colors = invert
        # 重置processed_image，需要重新处理
        self.processed_image = None
    
    def set_resize_options(self, enabled, width, height):
        """设置调整大小选项
        
        Args:
            enabled (bool): 是否启用调整大小
            width (int): 目标宽度
            height (int): 目标高度
        """
        self.resize_enabled = enabled
        self.resize_width = width
        self.resize_height = height
        # 重置processed_image，需要重新处理
        self.processed_image = None
        
    def set_rotate_options(self, enabled, angle):
        """设置旋转选项
        
        Args:
            enabled (bool): 是否启用旋转
            angle (int): 旋转角度，可以是0、90、180或270
        """
        self.rotate_enabled = enabled
        self.rotate_angle = angle
        # 重置processed_image，需要重新处理
        self.processed_image = None
        
    def set_flip_options(self, horizontal, vertical):
        """设置翻转选项
        
        Args:
            horizontal (bool): 是否水平翻转
            vertical (bool): 是否垂直翻转
        """
        self.flip_horizontal = horizontal
        self.flip_vertical = vertical
        # 重置processed_image，需要重新处理
        self.processed_image = None
    
    def apply_preprocess(self, image):
        """应用预处理（调整大小、旋转、翻转）
        
        Args:
            image (PIL.Image): 要处理的图像
            
        Returns:
            PIL.Image: 预处理后的图像
        """
        # 调整大小
        if self.resize_enabled:
            image = image.resize((self.resize_width, self.resize_height), Image.Resampling.LANCZOS)
            
        # 旋转
        if self.rotate_enabled:
            if self.rotate_angle == 90:
                image = image.transpose(Image.Transpose.ROTATE_90)
            elif self.rotate_angle == 180:
                image = image.transpose(Image.Transpose.ROTATE_180)
            elif self.rotate_angle == 270:
                image = image.transpose(Image.Transpose.ROTATE_270)
                
        # 翻转
        if self.flip_horizontal:
            image = ImageOps.mirror(image)
        if self.flip_vertical:
            image = ImageOps.flip(image)
            
        return image
    
    def process_image(self):
        """根据设置的阈值处理图像，将彩色图像转换为黑白像素图像
        
        Returns:
            PIL.Image: 处理后的图像，如果处理失败则返回None
        """
        if self.original_image is None:
            return None
            
        try:
            # 应用预处理
            img = self.apply_preprocess(self.original_image)
            
            # 转换为RGB模式确保可以获取RGB值
            img = img.convert('RGB')
            img_array = np.array(img)
            
            # 使用直接的阈值比较（0-255）
            # 只有当RGB三个通道都大于等于对应阈值时，才设为白色(255)，否则设为黑色(0)
            r_mask = img_array[:, :, 0] >= self.r_threshold
            g_mask = img_array[:, :, 1] >= self.g_threshold
            b_mask = img_array[:, :, 2] >= self.b_threshold
            
            # 合并掩码，所有条件都满足才为白色
            combined_mask = r_mask & g_mask & b_mask
            
            # 创建新的黑白图像数组
            bw_array = np.zeros_like(img_array)
            
            # 如果启用了颜色反转，则反转黑白
            if self.invert_colors:
                bw_array[~combined_mask] = 255
            else:
                bw_array[combined_mask] = 255
            
            # 将数组转换回PIL图像
            self.processed_image = Image.fromarray(bw_array.astype('uint8'))
            return self.processed_image
        except Exception as e:
            print(f"处理图像失败: {str(e)}")
            return None
            
    def process_single_image(self, image):
        """处理单个图像，用于批处理
        
        Args:
            image (PIL.Image): 要处理的图像
            
        Returns:
            PIL.Image: 处理后的图像，如果处理失败则返回None
        """
        try:
            # 应用预处理
            img = self.apply_preprocess(image)
            
            # 转换为RGB模式确保可以获取RGB值
            img = img.convert('RGB')
            img_array = np.array(img)
            
            # 使用直接的阈值比较（0-255）
            r_mask = img_array[:, :, 0] >= self.r_threshold
            g_mask = img_array[:, :, 1] >= self.g_threshold
            b_mask = img_array[:, :, 2] >= self.b_threshold
            
            # 合并掩码，所有条件都满足才为白色
            combined_mask = r_mask & g_mask & b_mask
            
            # 创建新的黑白图像数组
            bw_array = np.zeros_like(img_array)
            
            # 如果启用了颜色反转，则反转黑白
            if self.invert_colors:
                bw_array[~combined_mask] = 255
            else:
                bw_array[combined_mask] = 255
            
            # 将数组转换回PIL图像
            return Image.fromarray(bw_array.astype('uint8'))
        except Exception as e:
            print(f"处理图像失败: {str(e)}")
            return None
            
    def process_image_async(self):
        """异步处理图像，并发射进度信号
        
        Returns:
            bool: 处理启动成功返回True，否则返回False
        """
        if self.original_image is None:
            return False
            
        # 如果已经在处理中，则返回False
        if self.is_processing:
            return False
            
        # 标记为处理中
        self.is_processing = True
        
        # 创建线程进行异步处理
        self.processing_thread = threading.Thread(target=self._process_image_thread)
        self.processing_thread.daemon = True  # 设为守护线程
        self.processing_thread.start()
        
        return True
        
    def process_batch_async(self):
        """异步批量处理图像，并发射进度信号
        
        Returns:
            bool: 处理启动成功返回True，否则返回False
        """
        if not self.batch_images:
            return False
            
        # 如果已经在处理中，则返回False
        if self.is_processing:
            return False
            
        # 标记为处理中
        self.is_processing = True
        
        # 创建线程进行异步处理
        self.processing_thread = threading.Thread(target=self._process_batch_thread)
        self.processing_thread.daemon = True  # 设为守护线程
        self.processing_thread.start()
        
        return True
    
    def _process_image_thread(self):
        """图像处理线程函数"""
        try:
            # 发射初始进度信号
            self.signals.progress_changed.emit(0)
            
            # 应用预处理
            img = self.apply_preprocess(self.original_image)
            
            # 发射进度信号
            self.signals.progress_changed.emit(10)
            
            # 转换为RGB模式确保可以获取RGB值
            img = img.convert('RGB')
            img_array = np.array(img)
            
            # 发射进度信号
            self.signals.progress_changed.emit(20)
            
            # 使用直接的阈值比较（0-255）
            r_mask = img_array[:, :, 0] >= self.r_threshold
            
            # 发射进度信号
            self.signals.progress_changed.emit(40)
            
            g_mask = img_array[:, :, 1] >= self.g_threshold
            
            # 发射进度信号
            self.signals.progress_changed.emit(60)
            
            b_mask = img_array[:, :, 2] >= self.b_threshold
            
            # 合并掩码，所有条件都满足才为白色
            combined_mask = r_mask & g_mask & b_mask
            
            # 发射进度信号
            self.signals.progress_changed.emit(80)
            
            # 创建新的黑白图像数组
            bw_array = np.zeros_like(img_array)
            
            # 如果启用了颜色反转，则反转黑白
            if self.invert_colors:
                bw_array[~combined_mask] = 255
            else:
                bw_array[combined_mask] = 255
            
            # 将数组转换回PIL图像
            self.processed_image = Image.fromarray(bw_array.astype('uint8'))
            
            # 发射完成信号
            self.signals.progress_changed.emit(100)
            self.signals.processing_finished.emit(self.processed_image)
        except Exception as e:
            # 发射错误信号
            error_msg = f"处理图像失败: {str(e)}"
            print(error_msg)
            self.signals.processing_error.emit(error_msg)
        finally:
            # 标记为处理完成
            self.is_processing = False
            
    def _process_batch_thread(self):
        """批量处理线程函数"""
        try:
            # 清空结果列表
            self.batch_processed_images = []
            total_images = len(self.batch_images)
            
            # 逐个处理图像
            for i, image_data in enumerate(self.batch_images):
                # 处理当前图像
                img = image_data["image"]
                processed = self.process_single_image(img)
                
                if processed:
                    # 添加到结果列表
                    self.batch_processed_images.append({
                        "path": image_data["path"],
                        "image": processed,
                        "original": img
                    })
                else:
                    # 发射错误信号
                    self.signals.batch_processing_error.emit(
                        f"处理图像 {os.path.basename(image_data['path'])} 失败", i
                    )
            
            # 发射完成信号
            self.signals.batch_processing_finished.emit(self.batch_processed_images)
        except Exception as e:
            # 发射错误信号
            error_msg = f"批量处理失败: {str(e)}"
            print(error_msg)
            self.signals.processing_error.emit(error_msg)
        finally:
            # 标记为处理完成
            self.is_processing = False
    
    def get_processed_image(self):
        """获取处理后的图像，如果尚未处理则先进行处理
        
        Returns:
            PIL.Image: 处理后的图像
        """
        if self.processed_image is None:
            return self.process_image()
        return self.processed_image
        
    def is_busy(self):
        """检查处理器是否正在处理图像
        
        Returns:
            bool: 如果正在处理返回True，否则返回False
        """
        return self.is_processing
        
    def get_batch_results(self):
        """获取批处理结果
        
        Returns:
            list: 批处理结果列表
        """
        return self.batch_processed_images
    
    def export_bmp(self, output_path):
        """将处理后的图像导出为BMP格式，使用RGB565颜色模式
        
        Args:
            output_path (str): 输出文件路径
            
        Returns:
            bool: 导出成功返回True，否则返回False
        """
        if self.processed_image is None:
            self.process_image()
            
        if self.processed_image is None:
            return False
            
        try:
            # 获取基本文件名
            base_name = os.path.basename(output_path)
            dir_name = os.path.dirname(output_path)
            
            # 添加g_前缀
            if not base_name.startswith("g_"):
                new_name = f"g_{base_name}"
                output_path = os.path.join(dir_name, new_name)
            
            # 转换为RGB565格式
            # 注意：PIL不直接支持RGB565，所以将黑白图像保存为RGB模式
            self.processed_image.save(output_path, format='BMP')
            return True
        except Exception as e:
            print(f"导出BMP失败: {str(e)}")
            return False
    
    def export_batch_bmps(self, output_dir):
        """将批处理后的图像导出为BMP格式
        
        Args:
            output_dir (str): 输出目录路径
            
        Returns:
            dict: 导出结果，包含成功和失败的文件列表
        """
        if not self.batch_processed_images:
            return {"success": [], "failed": []}
            
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        result = {"success": [], "failed": []}
        
        # 逐个导出图像
        for image_data in self.batch_processed_images:
            try:
                # 获取基本文件名
                base_name = os.path.basename(image_data["path"])
                
                # 添加g_前缀
                if not base_name.startswith("g_"):
                    output_name = f"g_{base_name}"
                else:
                    output_name = base_name
                    
                # 构建完整输出路径
                output_path = os.path.join(output_dir, output_name)
                
                # 保存图像
                image_data["image"].save(output_path, format='BMP')
                
                # 添加到成功列表
                result["success"].append({
                    "input": image_data["path"],
                    "output": output_path
                })
            except Exception as e:
                print(f"导出BMP失败 {image_data['path']}: {str(e)}")
                # 添加到失败列表
                result["failed"].append({
                    "path": image_data["path"],
                    "error": str(e)
                })
        
        return result
        
    def export_batch_c_files(self, output_dir):
        """将批处理后的图像导出为C语言数组文件
        
        Args:
            output_dir (str): 输出目录路径
            
        Returns:
            dict: 导出结果，包含成功和失败的文件列表
        """
        if not self.batch_processed_images:
            return {"success": [], "failed": []}
            
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        result = {"success": [], "failed": []}
        
        # 逐个导出图像
        for image_data in self.batch_processed_images:
            try:
                # 获取基本文件名
                base_name = os.path.basename(image_data["path"])
                if base_name.lower().endswith('.bmp'):
                    base_name = base_name[:-4]
                
                # 构建完整输出路径
                output_path = os.path.join(output_dir, f"{base_name}.c")
                
                # 获取图像基本信息
                img = image_data["image"].convert('1')  # 转换为1位黑白图像
                width, height = img.size
                
                # 数组名使用gImage_前缀
                array_name = f"gImage_{base_name}"
                
                # 打开文件准备写入
                with open(output_path, 'w') as f:
                    # 写入数组声明
                    data_size = (width * height + 7) // 8  # 计算所需字节数，每8个像素占1个字节
                    f.write(f"const unsigned char {array_name}[{data_size}] = {{\n")
                    
                    # 处理像素数据
                    pixels = list(img.getdata())
                    byte_value = 0
                    bit_count = 0
                    byte_array = []
                    
                    # 将像素数据转换为字节数组
                    for pixel in pixels:
                        # 白色(255)表示1，黑色(0)表示0
                        bit_value = 1 if pixel > 0 else 0
                        byte_value = (byte_value << 1) | bit_value
                        bit_count += 1
                        
                        # 每8位组成一个字节
                        if bit_count == 8:
                            byte_array.append(byte_value)
                            byte_value = 0
                            bit_count = 0
                    
                    # 处理剩余的位
                    if bit_count > 0:
                        # 左移剩余的位
                        byte_value = byte_value << (8 - bit_count)
                        byte_array.append(byte_value)
                    
                    # 写入字节数据，每16个字节一行
                    for i, byte in enumerate(byte_array):
                        if i % 16 == 0 and i > 0:
                            f.write("\n")
                        f.write(f"0X{byte:02X}")
                        if i < len(byte_array) - 1:
                            f.write(", ")
                    
                    # 完成数组
                    f.write("\n};\n")
                
                # 添加到成功列表
                result["success"].append({
                    "input": image_data["path"],
                    "output": output_path
                })
            except Exception as e:
                print(f"导出C文件失败 {image_data['path']}: {str(e)}")
                # 添加到失败列表
                result["failed"].append({
                    "path": image_data["path"],
                    "error": str(e)
                })
        
        return result

    def export_c_file(self, output_path):
        """将处理后的图像导出为C语言数组文件
        
        Args:
            output_path (str): 输出文件路径
            
        Returns:
            bool: 导出成功返回True，否则返回False
        """
        if self.processed_image is None:
            self.process_image()
            
        if self.processed_image is None:
            return False
            
        try:
            # 获取图像基本信息
            img = self.processed_image.convert('1')  # 转换为1位黑白图像
            width, height = img.size
            
            # 获取基本文件名作为数组名前缀
            base_name = os.path.basename(output_path)
            if base_name.lower().endswith('.c'):
                base_name = base_name[:-2]
            
            # 数组名使用gImage_前缀
            array_name = f"gImage_{base_name}"
            
            # 打开文件准备写入
            with open(output_path, 'w') as f:
                # 写入数组声明
                data_size = (width * height + 7) // 8  # 计算所需字节数，每8个像素占1个字节
                f.write(f"const unsigned char {array_name}[{data_size}] = {{\n")
                
                # 处理像素数据
                pixels = list(img.getdata())
                byte_value = 0
                bit_count = 0
                byte_array = []
                
                # 将像素数据转换为字节数组
                for pixel in pixels:
                    # 白色(255)表示1，黑色(0)表示0
                    bit_value = 1 if pixel > 0 else 0
                    byte_value = (byte_value << 1) | bit_value
                    bit_count += 1
                    
                    # 每8位组成一个字节
                    if bit_count == 8:
                        byte_array.append(byte_value)
                        byte_value = 0
                        bit_count = 0
                
                # 处理剩余的位
                if bit_count > 0:
                    # 左移剩余的位
                    byte_value = byte_value << (8 - bit_count)
                    byte_array.append(byte_value)
                
                # 写入字节数据，每16个字节一行
                for i, byte in enumerate(byte_array):
                    if i % 16 == 0 and i > 0:
                        f.write("\n")
                    f.write(f"0X{byte:02X}")
                    if i < len(byte_array) - 1:
                        f.write(", ")
                
                # 完成数组
                f.write("\n};\n")
            
            return True
        except Exception as e:
            print(f"导出C文件失败: {str(e)}")
            return False 